import os
import sys
import torch
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.pyplot import MultipleLocator
import matplotlib.font_manager as fm

def get_project_root():
    return os.path.dirname(__file__)

def get_font_times(fontsize):
    return fm.FontProperties(
        fname=get_project_root() + 'fonts/times.ttf', size=fontsize)

def get_font_yahei(fontsize):
    return fm.FontProperties(
        fname=get_project_root() + 'fonts/msyh.ttf', size=fontsize)

def set_figsize(figsize=(30, 20)):
    """设置matplotlib的图表大小"""
    plt.rcParams['figure.figsize'] = figsize

def set_axes(axes, xlabel, ylabel, xlim, ylim, xscale, yscale, legend,
             xticker=None, yticker=None):
    """设置matplotlib的轴"""
    axes.set_xlabel(xlabel, fontproperties=get_font_times(40))
    axes.set_ylabel(ylabel, fontproperties=get_font_times(40))
    axes.set_xscale(xscale)
    axes.set_yscale(yscale)
    axes.set_xlim(xlim)
    axes.set_ylim(ylim)
    axes.tick_params(axis='both', which='major',labelsize=25)
    if not (xticker is None):
        axes.xaxis.set_major_locator(MultipleLocator(xticker))
    if not (yticker is None):
        axes.yaxis.set_major_locator(MultipleLocator(yticker))
    if legend:
        axes.legend(legend, fontsize=30)
    axes.grid()

class Animator:
    """在动画中绘制数据"""
    def __init__(self, xlabel=None, ylabel=None, legend=None, xlim=None,
                 ylim=None, xscale='linear', yscale='linear',
                 fmts=('-', 'm--', 'g-.', 'r:'), nrows=1, ncols=1,
                 figsize=(35, 25), title=None, xticker=None,yticker=None):
        # 增量地绘制多条线
        if legend is None:
            legend = []
        self.fig, self.axes = plt.subplots(nrows, ncols, figsize=figsize)
        if nrows * ncols == 1:
            self.axes = [self.axes, ]
        # 使用lambda函数捕获参数
        self.config_axes = lambda: set_axes(
            self.axes[0], xlabel, ylabel, xlim, ylim, xscale, yscale, legend, xticker,yticker)
        self.X, self.Y, self.fmts = None, None, fmts
        self.title = title
        plt.ion() # 使能动态更新图
        plt.show()

    def add(self, x, y):
        # 向图表中添加多个数据点
        if not hasattr(y, "__len__"):
            y = [y]
        n = len(y)
        if not hasattr(x, "__len__"):
            x = [x] * n
        if not self.X:
            self.X = [[] for _ in range(n)]
        if not self.Y:
            self.Y = [[] for _ in range(n)]
        for i, (a, b) in enumerate(zip(x, y)):
            if a is not None and b is not None:
                self.X[i].append(a)
                self.Y[i].append(b)
        self.axes[0].cla()
        for x, y, fmt in zip(self.X, self.Y, self.fmts):
            self.axes[0].plot(x, y, fmt)
        self.config_axes()
        self.axes[0].set_title(self.title, fontproperties=get_font_yahei(40))
        
        # 这两句使图可以更新
        self.fig.canvas.draw()
        self.fig.canvas.flush_events()
    
    def save_fig(self, fig_name):
        self.fig.savefig(f"{fig_name}.svg")

class Logger(object):
    def __init__(self, filename='default.log', stream=sys.stdout):
        self.terminal = stream
        self.log = open(filename, 'w')
 
    def write(self, message):
        self.terminal.write(message)
        self.log.write(message)
 
    def flush(self):
        # pass
        self.log.close()

class TextLogger:
    def __init__(self, log_path):
        self.log_path = log_path
        with open(self.log_path, "w") as f:
            f.write("")
    def log(self, log):
        with open(self.log_path, "a+") as f:
            f.write(log + "\n")
        print("LOG:\n", log)

# 规定情绪分类的标签名和数字标签
label_names = ["anger", "anxiety", "boredom", "disgust", "fear", "happiness", "sadness", "surprise", "uncertainty"]
label_nums = [0, 1, 2, 3, 4, 5, 6, 7, 8]

def label_name_to_num(label_name):
    return label_nums[label_names.index(label_name)]

def num_to_label_name(label_int):
    return label_names[label_nums.index(label_int)]

def get_all_subfolders(parent_folder):
    subfolder_List = []
    for file_name in os.listdir(parent_folder):
        full_path = os.path.join(parent_folder, file_name)
        if os.path.isdir(full_path):
            subfolder_List.append(full_path)
    subfolder_List.sort()
    return subfolder_List

def get_all_same_kind(parent_folder, extention_list):
    file_list = []
    path_list = os.listdir(parent_folder)
    for file_name in path_list:
        name, __fextention = os.path.splitext(file_name)
        if __fextention.lower() in extention_list:
            full_path = os.path.join(parent_folder, file_name)
            file_list.append(full_path)
    file_list.sort()
    return file_list

def print_cuda_info():
    print(torch.__version__)
    print(torch.version.cuda)
    print(torch.backends.cudnn.version())

def ensure_dir(path):
    os.makedirs(path, exist_ok=True)

"""
定义了一个方便转换不同数据模态路径的函数
传入任何一个模态的数据路径，可以得到对应样本另外的模态的数据路径
模态: video, pose_2d, pose_3d, pose_visual_2d, pose_visual_3d
"""
DATA_MODE = ["video",  "pose_2d", "pose_3d", "pose_2d_visual", "pose_3d_visual"]
data_suffix = ["", "_2d", "_3d", "_2d_vis", "_3d_vis"]

def emohugo_path_shift(path, from_data_mode, to_data_mode):
    assert from_data_mode in DATA_MODE
    assert to_data_mode in DATA_MODE
    assert from_data_mode != to_data_mode

    path, extention = os.path.splitext(path)
    path, data_name_mode = os.path.split(path)

    if from_data_mode == "video":
        path, category = os.path.split(path)
        path_root, data_mode_path = os.path.split(path)
        data_name = data_name_mode
    else:
        path, data_name = os.path.split(path)
        path, category = os.path.split(path)
        path_root, data_mode_path = os.path.split(path)

    data_name_mode = data_name + data_suffix[DATA_MODE.index(to_data_mode)]
    if to_data_mode == "video":
        data_mode_path = "emohugo_video"
        return os.path.join(path_root, data_mode_path, category, data_name + ".avi")
    elif to_data_mode == "pose_2d" or to_data_mode == "pose_3d":
        data_mode_path = "emohugo_pose"
        return os.path.join(path_root, data_mode_path, category, data_name, data_name_mode + ".json")
    elif to_data_mode == "pose_2d_visual" or to_data_mode == "pose_3d_visual":
        data_mode_path = "emohugo_pose_visual"
        return os.path.join(path_root, data_mode_path, category, data_name, data_name_mode + ".avi")

if __name__ == "__main__":
    print(get_project_root())

    json_path = os.path.join(get_project_root() , "datasets/emohugo_pose/disgust/000013-S001-15/000013-S001-15_2d.json")
    print("inferencing 3d from: ", json_path)

    reault_data_path = emohugo_path_shift(json_path, "pose_2d", "pose_3d")
    result_vis_path = emohugo_path_shift(json_path, "pose_2d", "pose_3d_visual")
    print("saving to: \n", reault_data_path, "\n", result_vis_path)

